﻿@page "/"
@using System.ComponentModel
@using Microsoft.Agents.AI
@inject IServiceProvider ServiceProvider
@inject NavigationManager Nav
@implements IDisposable

<PageTitle>Chat</PageTitle>

<ChatHeader OnNewChat="@ResetConversationAsync" />

<ChatMessageList Messages="@messages" InProgressMessage="@currentResponseMessage">
    <NoMessagesContent>
        <div>To get started, try asking about these example documents. You can replace these with your own data and replace this message.</div>
        <ChatCitation File="Example_Emergency_Survival_Kit.pdf"/>
        <ChatCitation File="Example_GPS_Watch.pdf"/>
    </NoMessagesContent>
</ChatMessageList>

<div class="chat-container">
    <ChatSuggestions OnSelected="@AddUserMessageAsync" @ref="@chatSuggestions" />
    <ChatInput OnSend="@AddUserMessageAsync" @ref="@chatInput" />
    <SurveyPrompt /> @* Remove this line to eliminate the template survey message *@
</div>

@code {
    private const string SystemPrompt = @"
        You are an assistant who answers questions about information you retrieve.
        Do not answer questions about anything else.
        Use only simple markdown to format your responses.

        Use the search tool to find relevant information. When you do this, end your
        reply with citations in the special XML format:

        <citation filename='string' page_number='number'>exact quote here</citation>

        Always include the citation in your response if there are results.

        The quote must be max 5 words, taken word-for-word from the search result, and is the basis for why the citation is relevant.
        Don't refer to the presence of citations; just emit these tags right at the end, with no surrounding text.
        ";

    private int statefulMessageCount;
    private readonly ChatOptions chatOptions = new();
    private readonly List<ChatMessage> messages = new();
    private CancellationTokenSource? currentResponseCancellation;
    private ChatMessage? currentResponseMessage;
    private ChatInput? chatInput;
    private ChatSuggestions? chatSuggestions;
    private AIAgent aiAgent = default!;

    protected override void OnInitialized()
    {
        // Resolve the keyed AI agent registered as "ChatAgent" in Program.cs
        aiAgent = ServiceProvider.GetRequiredKeyedService<AIAgent>("ChatAgent");
        
        statefulMessageCount = 0;
        messages.Add(new(ChatRole.System, SystemPrompt));
        //chatOptions.Tools = [AIFunctionFactory.Create(SearchAsync)];
    }

    private async Task AddUserMessageAsync(ChatMessage userMessage)
    {
        CancelAnyCurrentResponse();

        // Add the user message to the conversation
        messages.Add(userMessage);
        chatSuggestions?.Clear();
        await chatInput!.FocusAsync();

        // Stream and display a new response from the AI Agent
        var responseText = new TextContent("");
        currentResponseMessage = new ChatMessage(ChatRole.Assistant, [responseText]);
        currentResponseCancellation = new();
        
        // Use agent streaming instead of ChatClient streaming
        await foreach (var update in aiAgent.RunStreamingAsync(
            messages: messages.Skip(statefulMessageCount),
            cancellationToken: currentResponseCancellation.Token))
        {
            var responseUpdate = update.AsChatResponseUpdate();
            messages.AddMessages(responseUpdate, filter: c => c is not TextContent);
            responseText.Text += update.Text;
            chatOptions.ConversationId = responseUpdate.ConversationId;
            ChatMessageItem.NotifyChanged(currentResponseMessage);
        }

        // Store the final response in the conversation, and begin getting suggestions
        messages.Add(currentResponseMessage!);
        statefulMessageCount = chatOptions.ConversationId is not null ? messages.Count : 0;
        currentResponseMessage = null;
        chatSuggestions?.Update(messages);
    }

    private void CancelAnyCurrentResponse()
    {
        // If a response was cancelled while streaming, include it in the conversation so it's not lost
        if (currentResponseMessage is not null)
        {
            messages.Add(currentResponseMessage);
        }

        currentResponseCancellation?.Cancel();
        currentResponseMessage = null;
    }

    private async Task ResetConversationAsync()
    {
        CancelAnyCurrentResponse();
        messages.Clear();
        messages.Add(new(ChatRole.System, SystemPrompt));
        chatOptions.ConversationId = null;
        statefulMessageCount = 0;
        chatSuggestions?.Clear();
        await chatInput!.FocusAsync();
    }

    public void Dispose()
        => currentResponseCancellation?.Cancel();
}
